#ifndef _CUDA_MATRIX_H_
#define _CUDA_MATRIX_H_

/* matrix type of cuda matrix, this header can be included by cpp */
namespace cuda_matrix{

    /* matrix defined, row order storing */
    struct Mat{
        int y_max ,  x_max;
        size_t pitch;
        float *elem;

        Mat(){}
        Mat( int y_max, int x_max ){
            this->y_max = y_max;
            this->x_max = x_max; 
        } 
        inline void debug_info()const{
            printf("y_max=%d,x_max=%d\n",y_max,x_max);
        }
        inline size_t num_bytes() const{
            return y_max * pitch;
        }
        inline size_t line_width() const{
            return x_max * sizeof(float);
        }
        inline size_t num_line() const{
            return y_max;
        }
        
        inline float *operator[]( int y_idx ){
            return (float*)( (char*)elem + y_idx * pitch );
        }
        
        inline const float *operator[]( int y_idx )const{
            return (float*)( (char*)elem + y_idx * pitch );
        }
    };
    
    /* 3D matrix, an array of matrix in dense form */
    struct Mat3D{
        int z_max,y_max, x_max;
        size_t pitch;
        float *elem;
        Mat3D(){}
        Mat3D( int z_max, int y_max, int x_max ){
            this->z_max = z_max;
            this->y_max = y_max;
            this->x_max = x_max;
        }
        inline size_t num_bytes() const{
            return y_max * z_max * pitch;
        }
        inline size_t line_width() const{
            return x_max * sizeof(float);
        }
        inline size_t num_line() const{
            return y_max * z_max;
        }

        inline Mat operator[]( int z_idx ){
            Mat m;
            m.x_max = x_max;
            m.y_max = y_max;
            m.pitch = pitch;  
            m.elem = (float*)((char*)elem + z_idx * y_max * pitch ); 
            return m;
        }
        
        inline const Mat operator[]( int z_idx ) const{
            Mat m;
            m.x_max = x_max;
            m.y_max = y_max;
            m.pitch = pitch;  
            m.elem = (float*)((char*)elem + z_idx * y_max * pitch ); 
            return m;
        }                       
    };
    
    struct Mat4D{
        int v_max, h_max, y_max, x_max;
        size_t pitch;
        float *elem;
        Mat4D(){}
        Mat4D( int v_max, int h_max, int y_max, int x_max ){
            set_param( v_max, h_max, y_max, x_max );
        }

        inline void set_param( int v_max, int h_max, int y_max, int x_max ){
            this->v_max = v_max;
            this->h_max = h_max;
            this->y_max = y_max;
            this->x_max = x_max;
        }

        inline size_t num_bytes() const{
            return h_max * v_max * y_max * pitch;
        }
        inline size_t line_width() const{
            return x_max * sizeof(float);
        }
        inline size_t num_line() const{
            return y_max * h_max * v_max;
        }

        inline Mat3D operator[]( int v_idx ){
            Mat3D m;
            m.x_max = x_max;
            m.y_max = y_max;
            m.z_max = h_max;
            m.pitch = pitch;  
            m.elem = (float*)((char*)elem + v_idx * h_max * y_max * pitch ); 
            return m;
        }
        
        inline const Mat3D operator[]( int v_idx ) const{
            Mat3D m;
            m.x_max = x_max;
            m.y_max = y_max;
            m.z_max = h_max;
            m.pitch = pitch;  
            m.elem = (float*)((char*)elem + v_idx * h_max * y_max * pitch ); 
            return m;
        }                       
    };

    /* 
       4D matrix contains bias on each single 2D matrix, 
       this data structure is used specially for CRBM
    */
    struct Mat4DBias{
        Mat4D mat;
        float *v_bias, *h_bias;
        
        Mat4DBias(){}
        Mat4DBias( int v_max, int h_max, int y_max, int x_max ){
            mat.set_param( v_max, h_max, y_max, x_max );
        }
    };
};

#endif
